package stablediffusion

import (
	"app/pkg/config"
	fileTools "app/pkg/file"
	httpUtil "app/pkg/http"
	"app/pkg/log"
	"app/pkg/oss"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math"
)

type StableDiffusionClientV1 struct {
	BaseUrl []string
}

func NewStableDiffusionClient() *StableDiffusionClientV1 {
	client := &StableDiffusionClientV1{}
	client.WithBaseUrl(config.Config.Stablediffusion.Apiurl)
	return client
}

func (m *StableDiffusionClientV1) WithBaseUrl(baseUrl []string) {
	m.BaseUrl = baseUrl
}

func (m *StableDiffusionClientV1) CheckOptions(sdModelTitle string) (err error) {
	//是否切换 sd_model
	if curOptions, err := m.GetOptions(); err != nil {
		log.NewError("stable diffusion GetOptions Error", err.Error())
	} else if curOptions.SdModelCheckpoint != sdModelTitle {
		curOptions.SdModelCheckpoint = sdModelTitle
		if err := m.SetOptions(curOptions); err != nil {
			log.NewError("FindOne[model.SdModelModel] Error", err.Error())
		}
	}
	return
}

func (m *StableDiffusionClientV1) Txt2Img(data Txt2ImgRequest, ch chan int64) (pics []string, err error) {
	isEnd := false
	//这里更新进度
	go func() {
		for !isEnd {
			m.Progress(ch)
		}
		//这里来关闭ch 尽量发送端关闭channel原则
		close(ch)
	}()
	res, err := httpUtil.JsonPost(fmt.Sprintf("%s/sdapi/v1/txt2img", m.BaseUrl[0]), data, map[string]string{}, 60)
	isEnd = true
	if err != nil {
		return
	} else {
		fmt.Printf("%+v", string(res))
		txt2imgRes := Txt2imgResp{}
		err = json.Unmarshal(res, &txt2imgRes)
		if err != nil {
			return
		}
		//错误描述信息大于0
		if len(txt2imgRes.Detail) > 0 {
			err = errors.New(string(res))
			return
		}
		for _, v := range txt2imgRes.Images {
			bytes, _ := base64.StdEncoding.DecodeString(v)
			dstName := fileTools.DstName("img", "png")
			qnOss := oss.NewQnOss()
			//保存到本地
			err = fileTools.FilePutContent(string(bytes), fileTools.GetLocalPathByDstName(dstName))
			if err != nil {
				fmt.Println("FilePutContent: ", err)
			}
			//上传到阿里
			// pic := qnOss.PreviewUrl(dstName)
			if err := qnOss.PutImageBytes(dstName, bytes); err != nil {
				return pics, err
			}
			pics = append(pics, dstName)
		}
	}
	return
}

func (m *StableDiffusionClientV1) Progress(ch chan int64) {
	fmt.Println("progress query")
	res, _ := httpUtil.Get(fmt.Sprintf("%s/sdapi/v1/progress?skip_current_image=false", m.BaseUrl[0]), map[string]string{"Content-Type": "application/json; charset=utf-8"}, 60)
	proRes := ProgressResp{}
	json.Unmarshal(res, &proRes)
	if proRes.Progress != 0 {
		ch <- int64(math.Ceil(proRes.Progress * 100))
	}
}

// /sdapi/v1/options
func (m *StableDiffusionClientV1) SetOptions(options ConfigOptions) (err error) {
	_, err = httpUtil.JsonPost(fmt.Sprintf("%s/sdapi/v1/options", m.BaseUrl[0]), options, map[string]string{}, 60)
	return
}

func (m *StableDiffusionClientV1) ChangeCheckpoint(options ChangeCheckpointRequest) (err error) {
	if err := m.CheckpointAvariable(options.SdModelCheckpoint); err != nil {
		return err
	}
	_, err = httpUtil.JsonPost(fmt.Sprintf("%s/sdapi/v1/options", m.BaseUrl[0]), options, map[string]string{}, 60)
	return
}

func (m *StableDiffusionClientV1) CheckpointAvariable(ct string) (err error) {
	models, err := m.Models()
	if err != nil {
		return
	}
	for _, model := range models {
		if model.Title == ct {
			return
		}
	}
	err = fmt.Errorf("没有匹配的基础模型")
	return
}

func (m *StableDiffusionClientV1) GetOptions() (optins ConfigOptions, err error) {
	// fmt.Println(fmt.Sprintf("%s/sdapi/v1/options", m.BaseUrl[0]))
	res, err := httpUtil.Get(fmt.Sprintf("%s/sdapi/v1/options", m.BaseUrl[0]), map[string]string{}, 60)
	if err != nil {
		return
	} else {
		err = json.Unmarshal(res, &optins)
	}
	return
}

func (m *StableDiffusionClientV1) Models() (models []SdModelsReply, err error) {
	res, err := httpUtil.Get(fmt.Sprintf("%s/sdapi/v1/sd-models", m.BaseUrl[0]), map[string]string{}, 60)
	if err != nil {
		return
	} else {
		err = json.Unmarshal(res, &models)
	}
	return
}
